package wsz.rpc.zk;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.curator.framework.recipes.cache.NodeCache;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import wsz.rpc.common.ConstantsUtils;

/**
 * zk注册机制
 * @author wsz
 * @date 2021/12/7 18:12
 **/
public class ZkServer{
    private CuratorFramework client;
    // 节点名称(链接)、节点内容（）
    private String zkConnect;
    private String zkNodeName;

    private String currentPath;

    public ZkServer(String zkConnect, String zkNodeName) {
        try {
            if (zkConnect == null || "".equals(zkConnect)) {
                throw new RuntimeException("zkConnect cannot empty");
            }
            this.zkConnect = zkConnect;
            this.zkNodeName = zkNodeName;

            client = CuratorFrameworkFactory.builder()
                    .connectString(zkConnect)
                    .sessionTimeoutMs(50000)
                    .connectionTimeoutMs(3000)
                    .retryPolicy(new ExponentialBackoffRetry(1000, 3))
                    .build();

            client.start();

            registry();
        } catch (Exception ex) {
            System.out.println("cannot connect zk");
            ex.printStackTrace();
        }
    }

    public void closeServer() {
        if (client != null) {
            client.close();
        }
    }

    /**
     * 服务注册
     * @throws Exception
     */
    private void registry() throws Exception {
        initParentNode();
        initCurrentNode();
        bindListener();
    }

    /**
     * 节点监听器
     * @throws Exception
     */
    private void bindListener() throws Exception {
        final NodeCache nodeCache = new NodeCache(client, currentPath);
        nodeCache.getListenable().addListener(() -> {
            ChildData data = nodeCache.getCurrentData();
            if (data != null) {
                serverPrint("节点(上线)更新");
            } else {
                serverPrint("节点下线");
            }
        });
        nodeCache.start(true);
    }

    /**
     * 当前节点为：临时节点
     * @throws Exception
     */
    private void initCurrentNode() throws Exception {
        String path = ConstantsUtils.PARENT_NODE + "/" + zkNodeName;
        Stat stat = client.checkExists().forPath(path);
        if (stat == null) {
            currentPath = client.create()
                    .withMode(CreateMode.EPHEMERAL)
                    .forPath(path, "init".getBytes());
        } else {
            client.setData().forPath(path, "init".getBytes());
            currentPath = path;
        }
    }

    /**
     * 生成主节点：持久节点
     * @throws Exception
     */
    private void initParentNode() throws Exception {
        Stat stat = client.checkExists().forPath(ConstantsUtils.PARENT_NODE);
        if (stat == null) {
            client.create()
                    .withMode(CreateMode.PERSISTENT)
                    .forPath(ConstantsUtils.PARENT_NODE, ConstantsUtils.PARENT_NODE.getBytes());
        }
    }


    private void serverPrint(String msg) {
        System.out.println("zkServer " +msg +"：" + zkConnect + "_" + zkNodeName);
    }
}
